use thiserror::Error;

use stencil_tysys::EnvExpr;
use stencil_tysys::FnContext;
use stencil_tysys::{ArgRef, BinOp, FieldRef, FnContextExt, MonOp, ValueType};

use crate::{TaggedExpr, TypedExpr};

#[derive(Debug, Error)]
pub enum TypeError {
    #[error("arg {0:?} out of bounds")]
    ArgOutOfBounds(ArgRef),

    #[error("field {0:?} out of bounds or not in container context")]
    FieldOutOfBounds(FieldRef),

    #[error("invalid env reference {0:?}")]
    EnvRefUnknown(EnvExpr),

    #[error("invalid cast at expr {0} (from {1:?} to {2:?})")]
    InvalidCast(u32, ValueType, ValueType),

    #[error("type mismatch at expr {0} expected {1:?}")]
    TypeMismatch(u32, ValueType),

    #[error("cond expr {0} branches mismatch ({1:?} vs {2:?})")]
    CondMismatch(u32, ValueType, ValueType),

    #[error("incompatible operand at {0} for {1:?} ({2:?}")]
    MonopMismatch(u32, MonOp, ValueType),

    #[error("incompatible operands at {0} for {1:?} ({2:?}, {3:?})")]
    BinopMismatch(u32, BinOp, ValueType, ValueType),

    #[error("type {1:?} at expr {0} not indexable")]
    TypeNotIndexable(u32, ValueType),
}

/// Takes an expression tagged with identifiers and returns an expression tagged
/// with types.
pub fn typecheck_expr(
    expr: &TaggedExpr<u32>,
    ctx: &impl FnContext,
) -> Result<TypedExpr, TypeError> {
    let v = match expr {
        TaggedExpr::Const(_, v) => TaggedExpr::Const(v.ty(), *v),

        TaggedExpr::Env(_, e) => TaggedExpr::Env(
            ctx.resolve_env_expr_type(e)
                .ok_or(TypeError::EnvRefUnknown(*e))?,
            *e,
        ),

        TaggedExpr::Monop(id, op, a) => {
            let at = typecheck_expr(a, ctx)?;
            let ty = at.tag();

            match (op, ty) {
                (MonOp::Neg, ValueType::Bool) => TaggedExpr::Monop(ty.clone(), *op, Box::new(at)),
                (MonOp::LNot, ValueType::Bool) => TaggedExpr::Monop(ty.clone(), *op, Box::new(at)),

                _ => return Err(TypeError::MonopMismatch(*id, *op, ty.clone())),
            }
        }

        TaggedExpr::Binop(id, op, a, b) => {
            let at = typecheck_expr(a, ctx)?;
            let aty = at.tag();
            let bt = typecheck_expr(b, ctx)?;
            let bty = bt.tag();

            match (op, aty, bty) {
                (BinOp::Eq | BinOp::Ne, _, _) if aty == bty => {
                    TaggedExpr::Binop(ValueType::Bool, *op, Box::new(at), Box::new(bt))
                }

                (BinOp::Lt | BinOp::Le | BinOp::Gt | BinOp::Ge, _, _)
                    if aty == bty && aty.is_integer() && bty.is_integer() =>
                {
                    TaggedExpr::Binop(ValueType::Bool, *op, Box::new(at), Box::new(bt))
                }

                (BinOp::Add | BinOp::Sub | BinOp::Mul, _, _)
                    if aty == bty && aty.is_integer() && bty.is_integer() =>
                {
                    TaggedExpr::Binop(aty.clone(), *op, Box::new(at), Box::new(bt))
                }

                _ => return Err(TypeError::BinopMismatch(*id, *op, aty.clone(), bty.clone())),
            }
        }

        TaggedExpr::Cond(id, c, t, e) => {
            // Make sure the condition is a bool.
            let ct = typecheck_expr(c, ctx)?;
            if *ct.tag() != ValueType::Bool {
                return Err(TypeError::TypeMismatch(*c.tag(), ValueType::Bool));
            }

            // Compute the types of the two branches and make sure they're the same.
            let tt = typecheck_expr(t, ctx)?;
            let tty = tt.tag();

            let et = typecheck_expr(e, ctx)?;
            let ety = et.tag();
            if tty != ety {
                return Err(TypeError::CondMismatch(*id, tty.clone(), ety.clone()));
            }

            TaggedExpr::Cond(tty.clone(), Box::new(ct), Box::new(tt), Box::new(et))
        }

        TaggedExpr::Cast(id, target, a) => {
            let at = typecheck_expr(a, ctx)?;
            if !is_castable(at.tag(), target) {
                return Err(TypeError::InvalidCast(
                    *id,
                    at.tag().clone(),
                    target.clone(),
                ));
            }

            TaggedExpr::Cast(target.clone(), target.clone(), Box::new(at))
        }

        TaggedExpr::BufferLen(_, buf) => {
            let buft = typecheck_expr(buf, ctx)?;
            if !is_buf_readable(buft.tag()) {
                return Err(TypeError::TypeNotIndexable(*buf.tag(), buft.tag().clone()));
            }

            TaggedExpr::BufferLen(ValueType::USize, Box::new(buft))
        }

        TaggedExpr::BufferRead(_, buf, s, l) => {
            let buft = typecheck_expr(buf, ctx)?;
            if !is_buf_readable(buft.tag()) {
                return Err(TypeError::TypeNotIndexable(*buf.tag(), buft.tag().clone()));
            }

            let st = typecheck_expr(s, ctx)?;
            if *st.tag() != ValueType::USize {
                return Err(TypeError::TypeMismatch(*s.tag(), ValueType::USize));
            }

            let lt = typecheck_expr(l, ctx)?;
            if *lt.tag() != ValueType::USize {
                return Err(TypeError::TypeMismatch(*l.tag(), ValueType::USize));
            }

            TaggedExpr::BufferRead(
                ValueType::ByteSlice,
                Box::new(buft),
                Box::new(st),
                Box::new(lt),
            )
        }

        TaggedExpr::BufferReadInt(_, buf, off, ienc) => {
            let buft = typecheck_expr(buf, ctx)?;
            if !is_buf_readable(buft.tag()) {
                return Err(TypeError::TypeNotIndexable(*buf.tag(), buft.tag().clone()));
            }

            let offt = typecheck_expr(off, ctx)?;
            if *offt.tag() != ValueType::USize {
                return Err(TypeError::TypeMismatch(*off.tag(), ValueType::USize));
            }

            // Just resolve the integer down to an expression type.
            let real_ty = ienc.ty().into();
            TaggedExpr::BufferReadInt(real_ty, Box::new(buft), Box::new(offt), *ienc)
        }
    };

    Ok(v)
}

fn is_buf_readable(t: &ValueType) -> bool {
    match t {
        ValueType::ByteSlice | ValueType::ByteArray(_) => true,
        _ => false,
    }
}

fn is_castable(src: &ValueType, dest: &ValueType) -> bool {
    use ValueType::*;

    if src == dest {
        return true;
    }

    match (src, dest) {
        (Bool, U8)
        | (Bool, U64)
        | (U8, U64)
        | (U8, I64)
        | (U8, USize)
        | (U8, Bool)
        | (ByteArray(_), ByteSlice)
        | (ByteArray(1), U8)
        | (ByteArray(8), I64)
        | (ByteArray(8), U64) => true,
        _ => false,
    }
}
